import baostock as bs
import pandas as pd
import matplotlib.pyplot as plt
from utils.methods import get_pe, get_roe

END_DATE = '2025-09-25'
plt.rcParams['font.sans-serif'] = ['Noto Sans CJK', 'Microsoft YaHei', 'SimHei']  # Fallback fonts
plt.rcParams['axes.unicode_minus'] = False  # Fix minus sign display

def plot_pr(code):
    """
    绘制股票的PR图。
    - 需要登陆操作: baostock.login()
    :param code: 股票代码
    :return:
    """
    rs = bs.query_history_k_data_plus(code,
                                      "date",
                                      start_date='1990-11-26', end_date=END_DATE,
                                      frequency="m", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    data = []
    for index, row in result.iterrows():
        if row['date'][5:7] == '12':
            data.append({'date': row['date']})
    if len(data) <= 5:
        print(f'{code} 数据不足5年')
        return
    data = data[4:]
    for entry in data:
        try:
            date =  entry['date']
            pe = get_pe(code, date)
            year = int(date[:4])
            avg_roe = get_roe(code, year - 4, year)
            entry['pr'] = pe / (avg_roe * 100)
        except Exception as e:
            print(e)
            print(f'{code} 计算失败')
    print(data)
    data = list(filter(lambda x: 'pr' in x, data))
    dates = [entry['date'] for entry in data]
    prs = [entry['pr'] for entry in data]
    plt.figure(figsize=(10, 6))
    plt.plot(dates, prs, marker='o', linestyle='-', color='b')
    plt.title(f'{code} PR 图')
    plt.xlabel('日期')
    plt.ylabel('PR')
    plt.xticks(rotation=45)
    plt.grid(True)
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    code = input('请输入股票代码: ')
    bs.login()
    plot_pr(code=code)
    bs.logout()